package com.example.chat.agent.service.impl;

import com.example.chat.agent.service.ChromaDbService;
import com.example.chat.agent.service.StoreService;
import com.example.chat.common.constant.ChromaStoreContant;
import com.example.chat.common.exception.BusinessException;
import com.example.chat.common.utils.MyFileUtils;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.onnx.HuggingFaceTokenizer;
import dev.langchain4j.model.zhipu.ZhipuAiEmbeddingModel;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;
import java.util.Objects;

@Service
public class StoreServiceImpl implements StoreService {
    @Resource
    ChromaDbService chromaDbService;

    @Override
    public void loadFile(MultipartFile multipartFile) {
        if (multipartFile.isEmpty()) { throw new BusinessException(500, "未获取到文件"); }

        ChromaEmbeddingStore embeddingStore = chromaDbService.getEmbeddingStore(ChromaStoreContant.ZHIPU_TEXT);
        ZhipuAiEmbeddingModel embeddingModel = chromaDbService.getEmbeddingModel();

        embeddingStore.removeAll();

        Document document = MyFileUtils.loadFile(multipartFile, "F:\\" + multipartFile.getOriginalFilename());
        if (Objects.isNull(document)) { throw new BusinessException(500, "加载文档失败"); }

        DocumentSplitter splitter = DocumentSplitters.recursive(500, 20, new HuggingFaceTokenizer());
        List<TextSegment> segments = splitter.split(document);

        for (TextSegment segment : segments) {
            dev.langchain4j.data.embedding.Embedding content = embeddingModel.embed(segment).content();
            embeddingStore.add(content, segment);
        }
    }

}
